import os
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain_core.chat_history import (
    BaseChatMessageHistory,
    InMemoryChatMessageHistory,
)
from langchain_core.runnables.history import RunnableWithMessageHistory


ZHIPUAI_API_KEY = "118289aba4ea48e9a18559c26ee5121c.pHfVJtryCLwM6v4r"

store = {}

config = {"configurable": {"session_id": "abc2"}}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]



model = ChatOpenAI(
    model="glm-4",
    openai_api_key=ZHIPUAI_API_KEY,
    openai_api_base="https://open.bigmodel.cn/api/paas/v4/",
)

# response = model.invoke([HumanMessage(content="我是谁")])
# print(response)
with_message_history = RunnableWithMessageHistory(model, get_session_history)
# response = with_message_history.invoke([HumanMessage(content="Hi! I'm Bob")], config=config)

response = with_message_history.invoke([HumanMessage(content="What's my name")], config=config)


print(response)
